import boto3
from termcolor import colored
from datetime import datetime
import json
import sys
from pydoc import pipepager

author = {
    "name":"gl4ssesbo1",
    "twitter":"https://twitter.com/gl4ssesbo1",
    "github":"https://github.com/gl4ssesbo1",
    "blog":"https://www.pepperclipp.com/"
}

needs_creds = True

variables = {
    "SERVICE": {
        "value": "lambda",
        "required": "true",
        "description": "The service that will be used to run the module. It cannot be changed."
    },
    "FUNCTION-NAME": {
        "value": "",
        "required": "true",
        "description": "The name, ARN or partial ARN of a Function"
    },
    "PAYLOAD": {
        "value": "",
        "required": "false",
        "description": "The JSON that you want to provide to your Lambda function as input. Either set this or PAYLOAD-FILE."
    },
    "PAYLOAD-FILE": {
        "value": "",
        "required": "false",
        "description": "The file with the JSON that you want to provide to your Lambda function as input. Either set this or PAYLOAD."
    },
    "INVOCATION-TYPE": {
        "value": "",
        "required": "false",
        "description": "Either 'Event', 'RequestResponse' or 'DryRun'."
    }
}
description = "Invokes a Lambda function with InvokeFunction API."

aws_command = "aws lambda invoke --function-name <my-function> --payload <payload or file> -region <region> --profile <profile>"

def exploit(profile, workspace):
    now = datetime.now()
    dt_string = now.strftime("%d_%m_%Y_%H_%M_%S")
    file = "{}_ec2_enum_instances".format(dt_string)
    filename = "./workspaces/{}/{}".format(workspace, file)

    functionname = variables['FUNCTION-NAME']['value']
    payload = variables['PAYLOAD']['value']
    payloadfile = variables['PAYLOAD-FILE']['value']
    invovation = variables['INVOCATION-TYPE']['value']

    if payloadfile != "" and payload != "":
        print(
            colored("[*] Either set PAYLOAD-FILE or PAYLOAD.", "red")
        )

    else:
        try:
            if not payloadfile == "":
                with open(payloadfile, 'r') as pfile:
                    payload = json.load(pfile)

            response = profile.invoke(
                FunctionName=functionname,
                InvocationType=invovation,
                LogType='None',
                Payload=payload
            )

            print(colored("------------------------------------", "yellow", attrs=['bold']))
            print(colored("FunctionName: {}".format(functionname), "yellow", attrs=['bold']))
            print(colored("------------------------------------", "yellow", attrs=['bold']))

            for key, value in response.items():
                print(
                    "\t{}: {}".format(
                        colored(key, 'red'),
                        colored(value, 'blue')
                    )
                )
        except:
            e = sys.exc_info()
            print(colored("[*] {}".format(e), "red"))
